
from IPython.display import Audio, display, HTML
import librosa
import matplotlib.pyplot as plt
import numpy as np
import os


RATE = 22050

def forever():
    """Generator for iterating over integers from 0 onwards."""
    i = 0
    while True:
        yield i
        i += 1
        
def plot_wave(ax, wave, title='', range=None):
    ax.set_title(title)
    if range is not None:
        ax.plot(wave[range[0]:range[1]])
    else:
        ax.plot(wave)

def plot_mel(ax, mel, title=''):
    ax.set_title(title)
    pos = ax.imshow(mel.transpose(), aspect='auto', cmap='coolwarm', interpolation='nearest', origin='lower')
    # fig.colorbar(pos, orientation='horizontal')
    
def show_wave_and_mel(wave, mel, title=''):
    fig, axes = plt.subplots(2, dpi=100, figsize=(10, 5))
    plot_wave(axes[0], wave, title)
    n_zoom = int(0.5 * RATE)
    mid = len(wave) // 2
    # plot_wave(axes[1], wave, 'Zoomed in', (mid - n_zoom // 2, mid + n_zoom // 2))
    plot_mel(axes[1], mel, 'Mel spectrogram')
    plt.subplots_adjust(hspace=0.3)
    plt.show()
    display(Audio(wave, rate=RATE))

gid = 0
def show_results(name, title):
    display(HTML('<h1>%s</h1>' % title))
    for i in forever():
        try:
            if i == 0:
                wave = librosa.core.load(os.path.join(name, name + '.wav'))[0]
            else:
                wave = np.load(os.path.join(name, 'wave-%.5d.npy' % i))
            mel = np.load(os.path.join(name, 'mel-%.5d.npy' % i))
        except FileNotFoundError:
            break
        if i == 0:
            title = 'Original audio'
        else:
            title = ('Audio after %d WaveNet application' + ('' if i == 1 else 's')) % i
        show_wave_and_mel(wave, mel, title)
        
        global gid
        print('id: %d' % gid)
        gid += 1
        display(HTML('<hr style="height: 2px;">'))


show_results('chords', 'Chords')
show_results('tron', 'Speech from movie trailer')

id: 0

id: 1

id: 2

id: 3

id: 4

id: 5

id: 6

id: 7

id: 8

id: 9